#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date    : 2025-02-01-0:03
# @Author  : 
# @File    : test-txt.py
# @Software: PyCharm

import os
import json

classes = ['ph5', 'p26', 'pl40', 'pl60', 'pn', 'i5', 'p11', 'pne', 'pcl', 'pl50', 'pcr', 'w55', 'pl5', 'ph4.5',
           'pl80', 'pg', 'w28', 'w30', 'pl30', 'p19', 'i4l', 'i2r', 'pw3.2', 'pm20', 'pbp', 'p5', 'pl120', 'w24',
           'p13', 'w57', 'ip', 'p10', 'il100', 'il60', 'il90', 'pb', 'pl110', 'w59', 'il80', 'pl100', 'ph4', 'pmb',
           'p14', 'pl15', 'i4', 'p16', 'p3', 'pl70', 'pdd', 'pr70', 'w13', 'w32', 'i2', 'pr40', 'pm30', 'w63',
           'p12', 'p17', 'p18', 'im', 'pl20', 'p6', 'pw3.5', 'p27', 'pcd', 'i14', 'p2', 'p1', 'i12', 'wc', 'i10',
           'p23', 'w58', 'p25', 'ph3', 'pl90', 'pbm', 'w5', 'pl10', 'pss', 'pm55', 'phclr', 'i13', 'i1', 'ph2.2',
           'w47', 'pr60', 'w38', 'il50', 'w16', 'w22', 'p20', 'pn-2', 'iz', 'p9', 'p1n', 'ph4.3', 'ps', 'pm8', 'w3',
           'w21', 'p29', 'w18', 'pa10', 'pa14', 'pa13', 'il70', 'ph2', 'pr100', 'pr80', 'pm5', 'w45', 'pmr', 'w12',
           'ph2.9', 'pr50', 'il110', 'w42', 'p8', 'pt', 'pm35', 'pa12', 'w41', 'p28', 'ph3.5', 'pw4', 'pm2.5', 'w37',
           'ph5.3', 'ph5.5', 'ph2.8', 'i15', 'w10', 'pmblr', 'p21', 'ph4.2', 'pm15', 'pr30', 'pctl', 'w66', 'w46',
           'ph1.8', 'pm50', 'w20', 'w15', 'pl25', 'pm40', 'pa18', 'pa6', 'pw4.5', 'p15', 'ph2.5', 'p4', 'w35', 'pm10',
           'pr20', 'i3', 'ph3.2', 'pw3', 'ph2.4', 'ph4.8', 'pw4.2', 'phcs', 'ph2.1', 'w34', 'pc', 'pr45', 'pm2', 'pl35',
           'pcs', 'pw2.5', 'i11', 'w60', 'pr10', 'pa8', 'p24', 'w8', 'w14', 'pm13', 'pnlc', 'pclr', 'w56', 'w43',
           'ph3.8']


def convert(size, box):

    '''
    @size: (w, h)， 图片的高宽
    @box: (xmin, xmax, ymin, ymax), 标注框的坐标
    @return: (x_center, y_center, w2, h2), 返回目标中心坐标与相对高宽
    '''

    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def annos2txt(annos, out_dir):
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    h, w = 2048, 2048  # 图片大小
    image_name = annos["path"]  # eg test/10056.jpg
    out_txt = os.path.join(out_dir, image_name[:-4] + '.txt')
    with open(out_txt, 'a') as f:  # 新建对应 txt 文件

        for obj in annos['objects']:  # 遍历每个子列表即原 txt 的每一行
            cls = obj['category']
            if cls in classes:
                cls_id = classes.index(cls)  # 写入编号对应类别标签

                '''读取框的坐标'''
                xmin = obj['bbox']['xmin']
                ymin = obj['bbox']['ymin']
                xmax = obj['bbox']['xmax']
                ymax = obj['bbox']['ymax']
                bndbox = convert((w, h), (xmin, xmax, ymin, ymax))

                f.write(str(cls_id) + " " + " ".join([str(a) for a in bndbox]) + '\n')  # 写入


def main():
    filedir = "./annotations_all.json"  # json文件
    ids_file = "./test/ids.txt"  # 指定为 train/test
    ids = open(ids_file).read().splitlines()  # 获取 id 编号
    annos = json.loads(open(filedir).read())
    outdir = './'  # txt 保存目录

    img_num = len(ids)  # 统计当前图片数量
    cnt_id = 1  # 计数变量

    for imgid in ids:
        print('\rprocessing :[{} / {}]'.format(cnt_id, img_num), end="")
        cnt_id += 1

        # 跳过没有目标的图片
        if imgid in annos['imgs']:
            xml_cls = annos['imgs'][imgid]
            annos2txt(xml_cls, outdir)  # 生成 txt 文件
    ''' 
    annos['imgs'][imgid] 示例：
    {'path': 'test/10056.jpg', 
    'id': 10056, 
    'objects': [
    {'bbox': {'xmin': 452.475, 'ymin': 886.139, 'xmax': 468.3168, 'ymax': 916.8317}, 'category': 'i5'}, 
    {'bbox': {'xmin': 1274.26, 'ymin': 927.723, 'xmax': 1294.0594, 'ymax': 949.505}, 'category': 'i5'}, 
    {'bbox': {'xmin': 414.851, 'ymin': 877.228, 'xmax': 431.6832, 'ymax': 909.901}, 'category': 'pne'}, 
    {'bbox': {'xmin': 1215.84, 'ymin': 928.713, 'xmax': 1237.6237999999998, 'ymax': 950.495}, 'category': 'pne'}, 
    {'bbox': {'xmin': 2016.83, 'ymin': 910.891, 'xmax': 2032.6733, 'ymax': 934.6535}, 'category': 'pne'}]}
    '''


if __name__ == "__main__":
    main()
